# -*- coding:utf-8 -*-
################################################################################
# This script for
#
# Written by CathyQ
################################################################################

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
from config import Config, ResetConfig


def str2bool(v):
    if v.lower() in ['yes', 'true', 'y', 't', 1]:
        return True
    elif v.lower() in ['no', 'false', 'n', 'f', 0]:
        return False
    else:
        raise argparse.ArgumentTypeError("Unsupported value encountered.")


def main():
    args = Config()
    parser = argparse.ArgumentParser()
    parser.add_argument('--restart', type=str2bool, default=args.restart,
                        help=f'Restart to train the model [{args.restart}]')
    parser.add_argument('-s', '--scene', default=args.scene, choices=['Homograph', 'Dynamic', 'Static', 'Siftflow'],
                        help='Scenes choice in [Homograph/Dynamic/Static/Siftflow]')
    parser.add_argument('-b', '--batch_size_per_gpu', type=int, default=args.batch_per_gpu,
                        help=f'Batch size for each gpu [{args.batch_per_gpu}]')
    parser.add_argument('--loss_norm', default=args.loss_norm, choices=[1, 2],
                        help=f'Loss norm in [L1/ L2]')
    parser.add_argument('-gpu', '--num_GPUs', type=int, default=args.num_GPUs,
                        help=f'num_GPUs [{args.num_GPUs}]')
    parser.add_argument('-lr', '--learning_rate', type=float, default=args.learning_rate,
                        help=f'Learning rate [{args.learning_rate}]')
    parser.add_argument('-e', '--num_epochs', type=int, default=args.num_epochs,
                        help=f'# of epochs [{args.num_epochs}]')
    parser.add_argument('--sample_freq', type=int, default=args.sample_freq,
                        help=f'Sample frequent of epochs [2]')
    parser.add_argument('--test_freq', type=int, default=args.test_freq,
                        help=f'Test frequent of epochs [{args.test_freq}]')
    parser.add_argument('--view_freq', type=int, default=args.view_freq,
                        help=f'View frequent of epochs [{args.view_freq}]')
    parser.add_argument('-arch', '--architecture', default=args.architecture,
                        choices=['AHDR', "WuUnet", "WuResnet"],
                        help=f"Architecture in [{args.architecture}]")
    parser.add_argument('--mu', type=int, default=args.mu,
                        help=f'Tone mapping parameter mu [{args.mu}]')

    in_args = parser.parse_args()
    for key, item in vars(in_args).items():
        setattr(args, key, item)

    args = ResetConfig(args)

    return True

if __name__ == '__main__':
    main()
